"""
2.	对下面附件文件所示的网络输出数据，进行nms（非极大值抑制）处理，最后用图形显示nms处理前后的效果对比（共30分）
"""
# 题目要求：
# ①	显示nms处理前的所有框
# ②	显示nms处理后的所有框
# ③	nms算法处理结果正确
# ④	数据做合适的排序
# ⑤	处理逻辑正确
import matplotlib.pyplot as plt
import numpy as np


def sep(label = '', cnt=32):
    print('-' * cnt, label, '-' * cnt, sep='')


THRESH = 0.7
boxes = np.array([[100, 100, 210, 210, 0.72],
                  [280, 290, 420, 420, 0.8],
                  [220, 220, 320, 330, 0.92],
                  [105, 90, 220, 210, 0.71],
                  [230, 240, 325, 330, 0.81],
                  [305, 300, 420, 420, 0.9],
                  [215, 225, 305, 328, 0.6],
                  [150, 260, 290, 400, 0.99],
                  [102, 108, 208, 208, 0.72]])
x1 = boxes[:, 0]
y1 = boxes[:, 1]
x2 = boxes[:, 2]
y2 = boxes[:, 3]
scores = boxes[:, 4]
orders = scores.argsort()[::-1]
areas = (x2 - x1 + 1) * (y2 - y1 + 1)

idx_arr_to_keep = np.array([], dtype=np.int64)
idx_arr_rest = orders.copy()
eps = 1e-10
no = 0
while True:
    no += 1
    sep(no)
    if len(idx_arr_rest) == 0:
        break
    idx_keep = idx_arr_rest[0]
    idx_arr_to_keep = np.append(idx_arr_to_keep, [idx_keep])
    idx_arr_rest = idx_arr_rest[1:]
    if len(idx_arr_rest) == 0:
        break

    print('Before')
    print('keep', idx_arr_to_keep)
    print('rest', idx_arr_rest)

    x2_min = np.minimum(x2[idx_keep], x2[idx_arr_rest])
    x1_max = np.maximum(x1[idx_keep], x1[idx_arr_rest])
    y2_min = np.minimum(y2[idx_keep], y2[idx_arr_rest])
    y1_max = np.maximum(y1[idx_keep], y1[idx_arr_rest])
    xx = x2_min - x1_max + 1
    xx = np.maximum(xx, 0)
    yy = y2_min - y1_max + 1
    yy = np.maximum(yy, 0)
    I = xx * yy
    U = areas[idx_keep] + areas[idx_arr_rest] - I
    iou = I / (U + eps)
    idx_arr_rest = idx_arr_rest[iou < THRESH]

    print('After')
    print('keep', idx_arr_to_keep)
    print('rest', idx_arr_rest)


def show_boxes(boxes, idx):
    global spn
    spn += 1
    x1 = boxes[idx, 0]
    y1 = boxes[idx, 1]
    x2 = boxes[idx, 2]
    y2 = boxes[idx, 3]
    plt.subplot(spr, spc, spn)
    c = 'k'
    plt.plot([x1, x2], [y1, y1], color=c)
    plt.plot([x1, x2], [y2, y2], color=c)
    plt.plot([x1, x1], [y1, y2], color=c)
    plt.plot([x2, x2], [y1, y2], color=c)
    for i in idx:
        plt.annotate(f'#{i}', xy=[boxes[i, 0], boxes[i, 1]])


spr = 1
spc = 2
spn = 0
plt.figure(figsize=[13, 6])
show_boxes(boxes, tuple(range(len(boxes))))
print(idx_arr_to_keep)
show_boxes(boxes, idx_arr_to_keep)

plt.show()
